import base64
import math
import re
from io import BytesIO

import matplotlib.cm
import numpy as np
import torch
import torch.nn
from PIL import Image
import time
import os
import json

def compute_errors(gt, pred):
    pred = pred.detach().cpu().numpy()+1e-6
    gt = gt.detach().cpu().numpy()+1e-6
    thresh = np.maximum((gt / pred), (pred / gt))
    a1 = (thresh < 1.25).mean()
    a2 = (thresh < 1.25 ** 2).mean()
    a3 = (thresh < 1.25 ** 3).mean()
    rmse = (gt - pred) ** 2
    rmse = np.sqrt(rmse.mean())

    w1 = gt - gt.mean()
    w2 = pred - pred.mean()
    denom = np.sqrt(np.sum(w1 ** 2) * np.sum(w2 ** 2))
    if denom < 1e-10:
        zncc = 0
    else:
        zncc = np.sum(w1 * w2) / denom
    # rmse_log = (np.log(gt) - np.log(pred)) ** 2
    # rmse_log = np.sqrt(rmse_log.mean())
    #
    # err = np.log(pred) - np.log(gt)
    # silog = np.sqrt(np.mean(err ** 2) - np.mean(err) ** 2) * 100

    # log_10 = (np.abs(np.log10(gt) - np.log10(pred))).mean()

    return dict(rmse=rmse, zncc=zncc, a1=a1, a2=a2, a3=a3)


def compute_precision(gt, pred, datatype):
    pred = pred.detach().cpu().numpy()+1e-6
    gt = gt.detach().cpu().numpy()+1e-6
    e = np.abs(gt - pred)
    e = (20 - e) / (20 - 5)
    e[e <= 0] = 0
    e[e >= 1] = 1
    score = np.mean(e) * 100
    return score

class RunningAverage:
    def __init__(self):
        self.avg = 0
        self.count = 0

    def append(self, value):
        self.avg = (value + self.count * self.avg) / (self.count + 1)
        self.count += 1

    def get_value(self):
        return self.avg


class RunningAverageDict:
    def __init__(self):
        self._dict = None

    def update(self, new_dict):
        if self._dict is None:
            self._dict = dict()
            for key, value in new_dict.items():
                self._dict[key] = RunningAverage()

        for key, value in new_dict.items():
            self._dict[key].append(value)

    def get_value(self):
        return {key: value.get_value() for key, value in self._dict.items()}